package com.solver.demo.smt.main;

import org.sosy_lab.java_smt.SolverContextFactory;
import org.sosy_lab.java_smt.api.*;
import org.sosy_lab.java_smt.api.Model.ValueAssignment;

import java.math.BigInteger;
import java.util.HashSet;
import java.util.Set;

public class SmtModelExample {
    public static void main(String[] args) throws Exception {
        // 创建一个 Solver 上下文
        try (SolverContext context = SmtSolverUtil.solverContext()) {

            FormulaManager fmgr = context.getFormulaManager();
            BooleanFormulaManager bmgr = fmgr.getBooleanFormulaManager();
            IntegerFormulaManager imgr = fmgr.getIntegerFormulaManager();
            // 创建一些变量
            NumeralFormula.IntegerFormula x = imgr.makeVariable("x");
            NumeralFormula.IntegerFormula y = imgr.makeVariable("y");
            // 创建公式 x + y = 10
            BooleanFormula equation = imgr.equal(imgr.add(x, y), imgr.makeNumber(10));
            // 创建公式 x > 0 and y > 0  并且x<=6 y<=6
            BooleanFormula constraints = bmgr.and(
                    imgr.greaterThan(x, imgr.makeNumber(0)),
                    imgr.greaterThan(y, imgr.makeNumber(0)),
                    imgr.lessOrEquals(y, imgr.makeNumber(6)),
                    imgr.lessOrEquals(x, imgr.makeNumber(6)));
            // 合并所有约束
            BooleanFormula formula = bmgr.and(equation, constraints);
            Set<BooleanFormula> seenModels = new HashSet<>();
            int modelsFound = 0;
            int requiredModels = 5;
            try (ProverEnvironment prover = context.newProverEnvironment(
                    SolverContext.ProverOptions.GENERATE_MODELS)) {
                prover.addConstraint(formula);
                while (modelsFound < requiredModels && !prover.isUnsat()) {
                    try (Model model = prover.getModel()) {
                        System.out.println("Model " + (modelsFound + 1) + ":");
                        BooleanFormula modelConstraint = bmgr.makeFalse();
                        for (ValueAssignment assignment : model) {
                            System.out.println(assignment.getName() + " = " + assignment.getValue());
                            // Build a formula that represents the current model
                            modelConstraint = bmgr.or(modelConstraint,
                                    bmgr.not(imgr.equal((NumeralFormula.IntegerFormula) assignment.getKey(),
                                            imgr.makeNumber((BigInteger) assignment.getValue()))));
                        }
                        seenModels.add(modelConstraint);
                        modelsFound++;
                    }
                    // Add a constraint to exclude all seen models
                    for (BooleanFormula seenModel : seenModels) {
                        prover.addConstraint(seenModel);
                    }
                }
                if (modelsFound < requiredModels) {
                    System.out.println("Fewer than " + requiredModels + " models found.");
                }
            }
        }
    }
}
